local Class = require('lua.lang.Class')
local Instruction = require('lua.cfg.Instruction')
local Stack = require('lua.util.Stack')
local Lua = require('lua.vm.Lua')
local ControlFlowGraph = require('lua.cfg.ControlFlowGraph')
local BasicBlock = require('lua.cfg.BasicBlock')
local EdgeType = require('lua.cfg.EdgeType')

--- @class CfgState : Class
--- @field code number[]
--- @field instructions Instruction[]
--- @field taskStack Stack<function>
--- @field entersMap table<number, boolean>
local CfgState = Class {}

--- @param code number[]
--- @param filter fun(instruction: Instruction): Instruction
--- @return CfgState
function CfgState:new(code, filter)
	filter = filter or function(instruction)
		return instruction
	end
	--- @type CfgState
	local o = {}
	o.code = code
	local instructions = {}
	for i = 1, #code do
		instructions[i] = filter(Instruction:new(code[i]))
	end
	o.instructions = instructions
	o.taskStack = Stack:new()
	o.entersMap = {}
	return setmetatable(o, CfgState)
end

--- @param dest number
--- @return boolean
function CfgState:tryAppendEnter(dest)
	local instructions = self.instructions
	if instructions[dest] then
		self.entersMap[dest] = true
		return true
	end
	return false
end

--- @return void
function CfgState:parseConnection()
	local instructions = self.instructions
	for i = 1, #instructions do
		local instruction = instructions[i]
		if instruction:isLogic() or (instruction:getOpCode() ==
				Lua.OP_LOADBOOL and instruction:getArgC() == 0) then
			self:tryAppendEnter(i + 2)
		elseif instruction:isLoop() then
			if not self:tryAppendEnter(i + instruction:getArgSbx() + 1) then
				instruction:setArgSbx(0)
			end
		end
	end
end

--- @return ControlFlowGraph
function CfgState:analyze()
	self:parseConnection()
	local graph = ControlFlowGraph:new()
	self:addParseBlockTask(graph, 1, {}, function(block)
		graph.startBlock = block
	end)
	local stack = self.taskStack
	while stack:isNotEmpty() do
		stack:pop()()
	end
	return graph
end

--- @param graph ControlFlowGraph
--- @param start number
--- @param visited table<number, BasicBlock>
--- @param callback fun(block: BasicBlock): void
--- @return void
function CfgState:addParseBlockTask(graph, start, visited, callback)
	self.taskStack:push(function()
		self:parseBlock(graph, start, visited, callback)
	end)
end

--- @param graph ControlFlowGraph
--- @param start number
--- @param visited table<number, BasicBlock>
--- @param callback fun(block: BasicBlock): void
--- @return void
function CfgState:parseBlock(graph, start, visited, callback)
	local instructions = self.instructions
	local entersMap = self.entersMap
	local size = #instructions
	if start < 1 or start > size then
		if size == 0 then
			callback(BasicBlock:new())
		end
		return
	end
	local block = visited[start]
	if block then
		callback(block)
		return
	end
	block = BasicBlock:new()
	for pc = start, size do
		local instruction = instructions[pc]
		if entersMap[pc] ~= nil and pc ~= start then
			self:addParseBlockTask(graph, pc, visited, function(target)
				graph:addEdge(block, target, EdgeType.NEXT)
			end)
			break
		end
		assert(visited[pc] == nil, instruction:toString() .. ' is already parsed!')
		visited[pc] = block
		block:addInstruction(instruction)
		if instruction:isReturn() then
			break
		end
		if instruction:isLoop() then
			self:addParseBlockTask(graph, pc + instruction:getArgSbx() + 1, visited, function(jump)
				graph:addEdge(block, jump, EdgeType.JUMP)
				if instruction:getOpCode() == Lua.OP_JMP then
					return
				end
				self:addParseBlockTask(graph, pc + 1, visited, function(next)
					graph:addEdge(block, next, EdgeType.NEXT)
				end)
			end)
			break
		end
		if instruction:isLogic() or (instruction:getOpCode() == Lua.OP_LOADBOOL and instruction:getArgC() == 0) then
			self:addParseBlockTask(graph, pc + 2, visited, function(jump)
				graph:addEdge(block, jump, EdgeType.JUMP)
				self:addParseBlockTask(graph, pc + 1, visited, function(next)
					graph:addEdge(block, next, EdgeType.NEXT)
				end)
			end)
			break
		end
	end
	graph:addBlock(block)
	callback(block)
end

return CfgState